package com.adnaan.backend.websocket;

import com.adnaan.backend.entity.SystemMessage;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class SystemMessageWebSocketHandler extends TextWebSocketHandler {

    private final ConcurrentHashMap<Long, WebSocketSession> userSessions = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper;

    public SystemMessageWebSocketHandler(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        try {
            String query = session.getUri().getQuery();
            Long userId = Long.parseLong(query.split("=")[1]);
            userSessions.put(userId, session);
            log.info("用户 {} 建立WebSocket连接", userId);
        } catch (Exception e) {
            log.error("WebSocket连接建立失败", e);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        // 处理客户端消息，如果需要的话
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        userSessions.entrySet().removeIf(entry -> entry.getValue().equals(session));
        log.info("WebSocket连接关闭");
    }

    /**
     * 发送消息给指定用户
     */
    public void sendMessageToUser(SystemMessage message) {
        WebSocketSession session = userSessions.get(message.getReceiverId());
        if (session != null && session.isOpen()) {
            try {
                String jsonMessage = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(jsonMessage));
                log.info("发送消息给用户 {}: {}", message.getReceiverId(), jsonMessage);
            } catch (IOException e) {
                log.error("发送消息失败", e);
            }
        }
    }

    /**
     * 广播消息给所有在线用户
     */
    public void broadcastMessage(SystemMessage message) {
        String jsonMessage;
        try {
            jsonMessage = objectMapper.writeValueAsString(message);
        } catch (IOException e) {
            log.error("消息序列化失败", e);
            return;
        }

        userSessions.forEach((userId, session) -> {
            if (session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(jsonMessage));
                    log.info("广播消息给用户 {}: {}", userId, jsonMessage);
                } catch (IOException e) {
                    log.error("发送消息给用户 {} 失败", userId, e);
                }
            }
        });
    }

    /**
     * 获取当前在线用户数
     */
    public int getOnlineCount() {
        return userSessions.size();
    }
} 